//
// Created by hyh on 4/23/20.
//

#include <SimpleRecord.h>
#include <libavutil/avassert.h>

#define STREAM_FRAME_RATE 30
#define STREAM_PIX_FMT  AV_PIX_FMT_YUV420P
#define STREAM_DURATION   10.0
#define SCALE_FLAGS SWS_BICUBIC

static void add_stream(OutputStream *outputStream, AVFormatContext *formatContext,
                       AVCodec **codec,
                       enum AVCodecID codec_id);

static void open_video(AVFormatContext * avFormatCtx,AVCodec *avCodec,OutputStream *outputStream,AVDictionary *opt_arg);

static void open_audio(AVFormatContext *avFormatCtx, AVCodec *avCodec, OutputStream *outputStream, AVDictionary *opt_arg);

static int width;
static int height;

int write_buffer(void *opaque, uint8_t *buf, int buf_size){

    SimpleRecord * simpleRecord= (SimpleRecord *)opaque;

    static JNIEnv * jniEnv =NULL;

    if(jniEnv==NULL){
        if(simpleRecord->javaVM->AttachCurrentThread(&jniEnv,NULL)!=JNI_OK){
            LOGE("获取JNIEvn失败");
        }else{
            LOGE("获取JNIEvn成功");
        }
    }

    static jclass videoRecorder = NULL;
    if(videoRecorder==NULL){
        videoRecorder = jniEnv->GetObjectClass(simpleRecord->instance);
    }

    static jmethodID methodId = NULL;
    if(methodId==NULL){
        methodId = jniEnv->GetMethodID(videoRecorder, "onCallback", "([BI)V");
    }

    jbyteArray byteArray = jniEnv->NewByteArray(buf_size);
    jniEnv->SetByteArrayRegion(byteArray, 0, buf_size, (const jbyte *)buf);
    jniEnv->CallVoidMethod(simpleRecord->instance,methodId,byteArray,buf_size);
    jniEnv->DeleteLocalRef(byteArray);

//    FILE * fp_write= (FILE *)opaque;
//    int true_size=fwrite(buf,1,buf_size,fp_write);

    return buf_size;
}


void SimpleRecord::open(JavaVM *javaVM_,jobject instance_,const char *destUrl, int w, int h) {

    this->javaVM = javaVM_;
    this->instance = instance_;
    width = w;
    height = h;

    AVDictionary *opt = NULL;

    av_register_all();

    unsigned char *outBuffer = (unsigned char *)av_malloc(32768);

//    FILE * file = fopen(destUrl,"wb+");

    AVIOContext *avio_out = avio_alloc_context(outBuffer,32768,1,this,NULL,write_buffer,NULL);
    if(avio_out==NULL){
        LOGE("Could not alloc io context");
        return;
    }


    avformat_alloc_output_context2(&formatCtx,NULL,"mpegts",NULL);
    if(!formatCtx){
        LOGE("Could not deduce output format from file extension: using MPEG.\n");
        avformat_alloc_output_context2(&formatCtx, NULL, "mpeg", NULL);
    }

    if(!formatCtx){
        exit(1);
    }

    formatCtx->pb = avio_out;
    formatCtx->flags = AVFMT_FLAG_CUSTOM_IO;
    formatCtx->oformat->audio_codec = AV_CODEC_ID_AAC;
    formatCtx->oformat->video_codec = AV_CODEC_ID_H264;
    outputFormat = formatCtx->oformat;


    if(outputFormat->video_codec!=AV_CODEC_ID_NONE){
        add_stream(&videoOutputStream,formatCtx,&videoCodec,outputFormat->video_codec);
        have_video = 1;
    }

    if(outputFormat->video_codec!=AV_CODEC_ID_NONE){
        add_stream(&audioOutputStream,formatCtx,&audioCodec,outputFormat->audio_codec);
        have_audio = 1;
    }

    if(have_video){
        open_video(formatCtx,videoCodec,&videoOutputStream,opt);
    }

    if(have_audio){
        open_audio(formatCtx,audioCodec,&audioOutputStream,opt);
    }

//    av_dump_format(formatCtx,0,destUrl,1);
//
//    /* open the output file, if needed */
//    if (!(outputFormat->flags & AVFMT_NOFILE)) {
//        int ret = avio_open(&formatCtx->pb, destUrl, AVIO_FLAG_WRITE);
//        if (ret < 0) {
//            LOGE("Could not open '%s': %s\n", destUrl,
//                 av_err2str(ret));
//            return;
//        }
//    }

    int ret = avformat_write_header(formatCtx,&opt);
    if(ret<0){
        LOGE("Error occurred when opening output file: %s\n",
             av_err2str(ret));
        return;
    }

}



//void SimpleRecord::open(const char *destUrl, int w, int h) {
//
//    AVDictionary *opt = NULL;
////    av_dict_set(&opt, "movflags", "rtphint+faststart", 0);
//
//    av_register_all();
//
//    avformat_alloc_output_context2(&formatCtx,NULL,NULL,destUrl);
//    if(!formatCtx){
//        LOGE("Could not deduce output format from file extension: using MPEG.\n");
//        avformat_alloc_output_context2(&formatCtx, NULL, "mpeg", destUrl);
//    }
//
//    if(!formatCtx){
//        exit(1);
//    }
//
//    outputFormat = formatCtx->oformat;
//
//    if(outputFormat->video_codec!=AV_CODEC_ID_NONE){
//        add_stream(&videoOutputStream,formatCtx,&videoCodec,outputFormat->video_codec);
//        have_video = 1;
//    }
//
//    if(outputFormat->video_codec!=AV_CODEC_ID_NONE){
//        add_stream(&audioOutputStream,formatCtx,&audioCodec,outputFormat->audio_codec);
//        have_audio = 1;
//    }
//
//    if(have_video){
//        open_video(formatCtx,videoCodec,&videoOutputStream,opt);
//    }
//
//    if(have_audio){
//        open_audio(formatCtx,audioCodec,&audioOutputStream,opt);
//    }
//
//    av_dump_format(formatCtx,0,destUrl,1);
//
//    /* open the output file, if needed */
//    if (!(outputFormat->flags & AVFMT_NOFILE)) {
//        int ret = avio_open(&formatCtx->pb, destUrl, AVIO_FLAG_WRITE);
//        if (ret < 0) {
//            LOGE("Could not open '%s': %s\n", destUrl,
//                 av_err2str(ret));
//            return;
//        }
//    }
//
//    int ret = avformat_write_header(formatCtx,&opt);
//    if(ret<0){
//        LOGE("Error occurred when opening output file: %s\n",
//             av_err2str(ret));
//        return;
//    }
//}


static void add_stream(OutputStream *outputStream, AVFormatContext *formatContext,
                       AVCodec **codec,
                       enum AVCodecID codec_id){

    AVCodecContext *avCodecContext;
    *codec = avcodec_find_encoder(codec_id);
    if(!(*codec)){
        LOGE("Could not find encoder for '%s'\n",avcodec_get_name(codec_id));
        exit(1);
    }

    LOGE("Could not find encoder for '%s'\n",avcodec_get_name(codec_id));

    outputStream->st = avformat_new_stream(formatContext,NULL);
    if(!outputStream->st){
        LOGE("Could not allocate stream\n");
        exit(1);
    }

    outputStream->st->id = formatContext->nb_streams-1;
    LOGE("id=%d",outputStream->st->id);
    avCodecContext = avcodec_alloc_context3(*codec);
    if(!avCodecContext){
        LOGE("Could not alloc an encoding context\n");
        exit(1);
    }

    outputStream->avCodecCtx = avCodecContext;

    switch ((*codec)->type){
        case AVMEDIA_TYPE_AUDIO:
            avCodecContext->sample_fmt = (*codec)->sample_fmts?(*codec)->sample_fmts[0]:AV_SAMPLE_FMT_FLTP;
            avCodecContext->bit_rate = 64000;
            avCodecContext->sample_rate = 44100;
            if ((*codec)->supported_samplerates) {
                avCodecContext->sample_rate = (*codec)->supported_samplerates[0];
                for (int i = 0; (*codec)->supported_samplerates[i]; i++) {
                    if ((*codec)->supported_samplerates[i] == 44100)
                        avCodecContext->sample_rate = 44100;
                }
            }
            avCodecContext->channels = av_get_channel_layout_nb_channels(avCodecContext->channel_layout);
            avCodecContext->channel_layout = AV_CH_LAYOUT_STEREO;
            if ((*codec)->channel_layouts) {
                avCodecContext->channel_layout = (*codec)->channel_layouts[0];
                for (int i = 0; (*codec)->channel_layouts[i]; i++) {
                    if ((*codec)->channel_layouts[i] == AV_CH_LAYOUT_STEREO)
                        avCodecContext->channel_layout = AV_CH_LAYOUT_STEREO;
                }
            }
            avCodecContext->channels = av_get_channel_layout_nb_channels(avCodecContext->channel_layout);
            outputStream->st->time_base = (AVRational){1,avCodecContext->sample_rate};
            avCodecContext->time_base = outputStream->st->time_base;
            break;
        case AVMEDIA_TYPE_VIDEO:
            avCodecContext->codec_id = codec_id;
            avCodecContext->bit_rate = 400000;
            avCodecContext->width = width;
            avCodecContext->height = height;
            outputStream->st->time_base = (AVRational){1,STREAM_FRAME_RATE};
            avCodecContext->time_base = outputStream->st->time_base;
            avCodecContext->gop_size = 12;
            avCodecContext->pix_fmt = STREAM_PIX_FMT;
            if(avCodecContext->codec_id==AV_CODEC_ID_MPEG2VIDEO){
                avCodecContext->max_b_frames = 2;
            }
            if(avCodecContext->codec_id == AV_CODEC_ID_MPEG1VIDEO){
                avCodecContext->mb_decision = 2;
            }
            break;
    }

    if (formatContext->oformat->flags & AVFMT_GLOBALHEADER)
        avCodecContext->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;

}

static AVFrame *alloc_video_frame(enum AVPixelFormat pix_fmt, int width, int height);

static void open_video(AVFormatContext * avFormatCtx,AVCodec *avCodec,OutputStream *outputStream,AVDictionary *opt_arg){

    AVCodecContext * avCodecCtx = outputStream->avCodecCtx;
    AVDictionary *opt = NULL;
    av_dict_copy(&opt,opt_arg,0);
    // 设定H264编码参数，如果不设定，编码会有延迟
    if (avCodecCtx->codec_id == AV_CODEC_ID_H264) {
        av_dict_set(&opt, "tune", "zerolatency", 0);
        av_opt_set(avCodecCtx->priv_data, "preset", "ultrafast", 0);
        av_dict_set(&opt, "profile", "baseline", 0);
    }
    int ret = avcodec_open2(avCodecCtx,avCodec,&opt);
    av_dict_free(&opt);
    if(ret<0){
        LOGE("Could not open video codec: %s\n", av_err2str(ret));
        exit(1);
    }

    outputStream->frame = alloc_video_frame(avCodecCtx->pix_fmt,avCodecCtx->width,avCodecCtx->height);

    if (!outputStream->frame) {
        LOGE("Could not allocate video frame\n");
        exit(1);
    }

    /* If the output format is not YUV420P, then a temporary YUV420P
     * picture is needed too. It is then converted to the required
     * output format. */
    // 如果输出格式不是YUV420P，那么也需要临时的YUV420P图像。 然后将其转换为所需的输出格式
    outputStream->tmp_frame = NULL;
    if (avCodecCtx->pix_fmt != AV_PIX_FMT_YUV420P) {
        outputStream->tmp_frame = alloc_video_frame(AV_PIX_FMT_YUV420P, avCodecCtx->width, avCodecCtx->height);
        if (!outputStream->tmp_frame) {
            LOGE("Could not allocate temporary picture\n");
            exit(1);
        }
    }

    /* copy the stream parameters to the muxer */
    ret = avcodec_parameters_from_context(outputStream->st->codecpar, avCodecCtx);
    if (ret < 0) {
        LOGE("Could not copy the stream parameters\n");
        exit(1);
    }
}

static AVFrame *alloc_audio_frame(enum AVSampleFormat sample_fmt,
                                  uint64_t channel_layout,
                                  int sample_rate, int nb_samples);

static void open_audio(AVFormatContext *avFormatCtx, AVCodec *avCodec, OutputStream *outputStream, AVDictionary *opt_arg){


    AVCodecContext *avCodecCtx = outputStream->avCodecCtx;
    AVDictionary *opt = NULL;

    av_dict_copy(&opt,opt_arg,0);
    int ret = avcodec_open2(avCodecCtx,avCodec,&opt);
    av_dict_free(&opt);
    if (ret < 0) {
        LOGE("Could not open audio codec: %s\n", av_err2str(ret));
        exit(1);
    }


    int nb_samples;
    if(avCodecCtx->codec->capabilities&AV_CODEC_CAP_VARIABLE_FRAME_SIZE){
        nb_samples = 10000;
    }else{
        nb_samples = avCodecCtx->frame_size;
    }

    LOGE("nb_samples=%d",nb_samples);
    LOGE("sample_fmt=%s",av_get_sample_fmt_name(avCodecCtx->sample_fmt));
    outputStream->frame = alloc_audio_frame(avCodecCtx->sample_fmt,avCodecCtx->channel_layout,avCodecCtx->sample_rate,nb_samples);

    outputStream->tmp_frame = alloc_audio_frame(AV_SAMPLE_FMT_S16P,avCodecCtx->channel_layout,avCodecCtx->sample_rate,nb_samples);

    /* copy the stream parameters to the muxer */
    ret = avcodec_parameters_from_context(outputStream->st->codecpar, avCodecCtx);
    if (ret < 0) {
        LOGE("Could not copy the stream parameters\n");
        exit(1);
    }

    outputStream->swr_ctx = swr_alloc();
    if(!outputStream->swr_ctx){
        LOGE("Could not allocate resampler context\n");
        exit(1);
    }

    av_opt_set_int       (outputStream->swr_ctx, "in_channel_count",   avCodecCtx->channels,       0);
    av_opt_set_int       (outputStream->swr_ctx, "in_sample_rate",     avCodecCtx->sample_rate,    0);
    av_opt_set_sample_fmt(outputStream->swr_ctx, "in_sample_fmt",      AV_SAMPLE_FMT_S16, 0);
    av_opt_set_int       (outputStream->swr_ctx, "out_channel_count",  avCodecCtx->channels,       0);
    av_opt_set_int       (outputStream->swr_ctx, "out_sample_rate",    avCodecCtx->sample_rate,    0);
    av_opt_set_sample_fmt(outputStream->swr_ctx, "out_sample_fmt",     avCodecCtx->sample_fmt,     0);

    ret=swr_init(outputStream->swr_ctx);
    if(ret<0){
        LOGE("Failed to initialize the resampling context\n");
        exit(1);
    }

}



static AVFrame *alloc_video_frame(enum AVPixelFormat pix_fmt, int width, int height)
{
    AVFrame *picture;
    int ret;
    picture = av_frame_alloc();
    if (!picture)
        return NULL;
    picture->format = pix_fmt;
    picture->width  = width;
    picture->height = height;
    /* allocate the buffers for the frame data */
    ret = av_frame_get_buffer(picture, 32);
    if (ret < 0) {
        LOGE("Could not allocate frame data.\n");
        exit(1);
    }
    return picture;
}

static AVFrame *alloc_audio_frame(enum AVSampleFormat sample_fmt,
                                  uint64_t channel_layout,
                                  int sample_rate, int nb_samples)
{
    AVFrame *frame = av_frame_alloc();
    int ret;
    if (!frame) {
        fprintf(stderr, "Error allocating an audio frame\n");
        exit(1);
    }
    frame->format = sample_fmt;
    frame->channel_layout = channel_layout;
    frame->sample_rate = sample_rate;
    frame->nb_samples = nb_samples;
    if (nb_samples) {
        ret = av_frame_get_buffer(frame, 0);
        LOGE("ret=%d",ret);
        if (ret < 0) {
            fprintf(stderr, "Error allocating an audio buffer\n");
            exit(1);
        }
    }
    return frame;
}


static AVFrame *get_audio_frame(OutputStream *ost,u_int8_t *data,int len);

static AVFrame *get_video_frame(OutputStream *ost,u_int8_t *data,int len);

static int write_frame(AVFormatContext *fmt_ctx, const AVRational *time_base, AVStream *st, AVPacket *pkt);

void SimpleRecord::encodeVideo(uint8_t *data, int len) {


    AVPacket pkt = {0};
    AVCodecContext *avCodecCtx = videoOutputStream.avCodecCtx;

    AVFrame * frame = get_video_frame(&videoOutputStream,data,len);

    av_init_packet(&pkt);

    int got_packet = 0;

    int ret = avcodec_encode_video2(avCodecCtx,&pkt,frame,&got_packet);

    if (ret < 0) {
        LOGE("Error encoding video frame: %s\n", av_err2str(ret));
        exit(1);
    }
    if (got_packet) {
        ret = write_frame(formatCtx, &avCodecCtx->time_base, videoOutputStream.st, &pkt);
    } else {
        ret = 0;
    }
    av_free_packet(&pkt);
    if (ret < 0) {
        LOGE("Error while writing video frame: %s\n", av_err2str(ret));
        exit(1);
    }
}

void SimpleRecord::encodeAudio(uint8_t *data, int len) {

    AVPacket pkt = {0};
    av_init_packet(&pkt);
    AVCodecContext * avCodecContext = audioOutputStream.avCodecCtx;
    AVFrame * frame = get_audio_frame(&audioOutputStream,data,len);
    if(frame){
        /* convert samples from native format to destination codec format, using the resampler */
        /* compute destination number of samples */
        int dst_nb_samples = av_rescale_rnd(swr_get_delay(audioOutputStream.swr_ctx, avCodecContext->sample_rate) + frame->nb_samples,
                                            avCodecContext->sample_rate, avCodecContext->sample_rate, AV_ROUND_UP);
        av_assert0(dst_nb_samples == frame->nb_samples);
        /* when we pass a frame to the encoder, it may keep a reference to it
         * internally;
         * make sure we do not overwrite it here
         */
        int ret = av_frame_make_writable(audioOutputStream.frame);
        if (ret < 0)
            exit(1);
        /* convert to destination format */
        ret = swr_convert(audioOutputStream.swr_ctx,
                          audioOutputStream.frame->data, dst_nb_samples,
                          (const uint8_t **)frame->data, frame->nb_samples);
        if (ret < 0) {
            LOGE("Error while converting\n");
            exit(1);
        }

        frame = audioOutputStream.frame;

        frame->pts = av_rescale_q(audioOutputStream.samples_count, (AVRational){1, avCodecContext->sample_rate}, avCodecContext->time_base);
//        LOGE("dst_nb_samples:%d pts:%lld",dst_nb_samples,frame->pts);
        audioOutputStream.samples_count += dst_nb_samples;
    }
    int got_packet;
    int ret = avcodec_encode_audio2(avCodecContext, &pkt, frame, &got_packet);
    if (ret < 0) {
        LOGE("Error encoding audio frame: %s\n", av_err2str(ret));
        exit(1);
    }

    if (got_packet) {
        ret = write_frame(formatCtx, &avCodecContext->time_base, audioOutputStream.st, &pkt);
        if (ret < 0) {
            LOGE("Error while writing audio frame: %s\n",
                 av_err2str(ret));
            exit(1);
        }
    }
    av_free_packet(&pkt);
}


static void fill_yuv_image(AVCodecContext *c, AVFrame *frame, int frame_index,
                           int width, int height,u_int8_t *data);

static AVFrame *get_video_frame(OutputStream *ost,u_int8_t *data,int len)
{
    AVCodecContext *c = ost->avCodecCtx;
    /* check if we want to generate more frames */
//    if (av_compare_ts(ost->next_pts, c->time_base,
//                      STREAM_DURATION, (AVRational){ 1, 1 }) >= 0){
//        LOGE("视频丢一帧");
//        return NULL;
//    }
    /* when we pass a frame to the encoder, it may keep a reference to it
     * internally; make sure we do not overwrite it here */
    if (av_frame_make_writable(ost->frame) < 0)
        exit(1);
    if (c->pix_fmt != AV_PIX_FMT_YUV420P) {
        /* as we only generate a YUV420P picture, we must convert it
         * to the codec pixel format if needed */
        if (!ost->sws_ctx) {
            ost->sws_ctx = sws_getContext(c->width, c->height,
                                          AV_PIX_FMT_YUV420P,
                                          c->width, c->height,
                                          c->pix_fmt,
                                          SCALE_FLAGS, NULL, NULL, NULL);
            if (!ost->sws_ctx) {
                LOGE("Could not initialize the conversion context\n");
                exit(1);
            }
        }
        fill_yuv_image(c,ost->tmp_frame, ost->next_pts, c->width, c->height,data);

        sws_scale(ost->sws_ctx,
                  (const uint8_t * const *)ost->tmp_frame->data, ost->tmp_frame->linesize,
                  0, c->height, ost->frame->data, ost->frame->linesize);
    } else {
        fill_yuv_image(c,ost->frame, ost->next_pts, c->width, c->height,data);
    }
    ost->frame->pts = ost->next_pts++;
    return ost->frame;
}

static AVFrame *get_audio_frame(OutputStream *ost,u_int8_t *data,int len)
{
    /* check if we want to generate more frames */
//    if (av_compare_ts(ost->next_pts, ost->avCodecCtx->time_base,
//                      STREAM_DURATION, (AVRational){ 1, 1 }) >= 0){
//        LOGE("音频丢一帧");
//        return NULL;
//    }
//    LOGE("len=%d",len);
    unsigned char *srcData = data;
    AVFrame *frame = ost->tmp_frame;
    memcpy(frame->data[0],srcData,len);
    frame->pts = ost->next_pts;
    ost->next_pts+=frame->nb_samples;
    return frame;
}




/* Prepare a dummy image. */
static void fill_yuv_image(AVCodecContext *c, AVFrame *frame, int frame_index,
                           int width, int height,u_int8_t *data)
{
    int size = width*height;
    memcpy(frame->data[0],data,size);
    for (int j = 0; j < size / 4; j++) {
        *(frame->data[2] + j) = *(data + size + j * 2);
        // V
        *(frame->data[1] + j) = *(data + size + j * 2 + 1);
        //U
    }
}


static void log_packet(const AVFormatContext *fmt_ctx, const AVPacket *pkt);

mutex writeMutex;
static int write_frame(AVFormatContext *fmt_ctx, const AVRational *time_base, AVStream *st, AVPacket *pkt)
{
    unique_lock<mutex> lock(writeMutex);
    /* rescale output packet timestamp values from codec to stream timebase */
    av_packet_rescale_ts(pkt, *time_base, st->time_base);
    if(st->index==1){
//        LOGE("pts=%lld num=%d num=%d",pkt->pts,time_base->den,st->time_base.den);
    }
    pkt->stream_index = st->index;
    /* Write the compressed frame to the media file. */
    log_packet(fmt_ctx, pkt);
    return av_interleaved_write_frame(fmt_ctx, pkt);
}

static void log_packet(const AVFormatContext *fmt_ctx, const AVPacket *pkt)
{
    AVRational *time_base = &fmt_ctx->streams[pkt->stream_index]->time_base;
//    LOGE("pts:%s pts_time:%s dts:%s dts_time:%s duration:%s duration_time:%s stream_index:%d\n",
//           av_ts2str(pkt->pts), av_ts2timestr(pkt->pts, time_base),
//           av_ts2str(pkt->dts), av_ts2timestr(pkt->dts, time_base),
//           av_ts2str(pkt->duration), av_ts2timestr(pkt->duration, time_base),
//           pkt->stream_index);
}

static void close_stream(AVFormatContext *oc, OutputStream *ost)
{
    avcodec_free_context(&ost->avCodecCtx);
    av_frame_free(&ost->frame);
    av_frame_free(&ost->tmp_frame);
    sws_freeContext(ost->sws_ctx);
    swr_free(&ost->swr_ctx);
}

void SimpleRecord::stop() {

    av_write_trailer(formatCtx);
    /* Close each codec. */
    if (have_video)
        close_stream(formatCtx, &videoOutputStream);
    if (have_audio)
        close_stream(formatCtx, &audioOutputStream);
    if (!(outputFormat->flags & AVFMT_NOFILE))
        /* Close the output file. */
        avio_closep(&formatCtx->pb);
    /* free the stream */
    avformat_free_context(formatCtx);
}


